package com.zon.len.mp.extend.configuration;

import com.baomidou.mybatisplus.core.conditions.AbstractLambdaWrapper;
import com.baomidou.mybatisplus.core.mapper.BaseMapper;
import com.baomidou.mybatisplus.core.mapper.Mapper;
import com.baomidou.mybatisplus.core.toolkit.Constants;
import com.baomidou.mybatisplus.core.toolkit.ReflectionKit;
import com.baomidou.mybatisplus.core.toolkit.StringPool;
import com.zon.len.mp.extend.matedata.TableMetadataHelper;
import com.zon.len.mp.extend.script.ExtendMybatisCommand;
import com.zon.len.mp.extend.wrapper.AbstractExtendWrapper;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import org.apache.ibatis.binding.MapperMethod.ParamMap;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

/**
 * 分库分表插件
 *
 * @author ZonLen since on 2021/10/8 下午12:55
 */
@Intercepts(
    {
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class,
            Object.class, RowBounds.class, ResultHandler.class}),
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class,
            Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class}),
    }
)
public class MybatisPlusShardingInterceptor implements Interceptor {

  private static final Map<String, Class<?>> MAPPED_STATEMENT_TABLE_CLASS = new ConcurrentHashMap<>();

  @Override
  public Object intercept(Invocation invocation) throws Throwable {
    final Object[] args = invocation.getArgs();
    final MappedStatement ms = (MappedStatement) args[0];
    doMpSelfTableSharding(ms, args[1]);
    return invocation.proceed();
  }

  private void doMpSelfTableSharding(MappedStatement ms, Object parameter) throws Throwable {
    final String statementId = ms.getId();
    final String methodName = statementId.substring(ms.getId().lastIndexOf(StringPool.DOT) + 1);
    final Class<?> aClass = MAPPED_STATEMENT_TABLE_CLASS.get(statementId);
    if (null != aClass) {
      final AbstractLambdaWrapper<?, ?> wrapper = (AbstractLambdaWrapper<?, ?>) ((ParamMap<?>) parameter)
          .get(Constants.WRAPPER);
      bindingTableNameParam(wrapper, aClass, methodName);
      return;
    }
    final boolean extendStatement = ExtendMybatisCommand.isExtendStatement(methodName);
    if (!extendStatement || parameter == null) {
      return;
    }
    if (parameter instanceof ParamMap) {
      ParamMap<?> param = (ParamMap<?>) parameter;
      final Object o = param.get(Constants.WRAPPER);
      if (o instanceof AbstractLambdaWrapper) {
        final String mapperClassStr = ms.getId()
            .substring(0, ms.getId().lastIndexOf(StringPool.DOT));
        final Class<?> mapperClass = Class.forName(mapperClassStr);
        if (!BaseMapper.class.isAssignableFrom(mapperClass)) {
          return;
        }
        final Class<?> tableClass = ReflectionKit
            .getSuperClassGenericType(mapperClass, Mapper.class, 0);
        MAPPED_STATEMENT_TABLE_CLASS.putIfAbsent(statementId, tableClass);
        bindingTableNameParam((AbstractLambdaWrapper<?, ?>) o, tableClass, methodName);
      }
    }
  }

  private void bindingTableNameParam(AbstractLambdaWrapper<?, ?> wrapper, Class<?> tableClass,
      String methodName) {
    String tableName;
    if (wrapper instanceof AbstractExtendWrapper) {
      tableName = ((AbstractExtendWrapper<?, ?>) wrapper).getFinalTableName();
    } else {
      tableName = TableMetadataHelper.globalFinalTableName(tableClass);
    }
    wrapper.getParamNameValuePairs().put(TableMetadataHelper.tableMetadata(tableClass).getAlias(),
        tableName);
  }

  @Override
  public Object plugin(Object target) {
    if (target instanceof Executor) {
      return Plugin.wrap(target, this);
    }
    return target;
  }

}
